

import pandas as pd
import sys

#%% SET PATH ----------------------------------------------------------------------------
# When running this file usually also add argument which is the path for the table
try:
    overleaf_path = sys.argv[1]
except:
    overleaf_path = './../Apps/Overleaf/bbm/draft/reports/revision/'


for k in ['boom_bust', '2005']:
    df_contracts = pd.read_csv('./data_py/processed/contracts_final.csv', index_col=[0])

    #%% GET PRICE RANKING BY MONTH SET PATH ---------------------------------------------
    # Only look at rigs with at least one contract before/after

    if k == 'boom_bust':
        df_contracts['first_5'] = df_contracts['boom']
    elif k == '2005':
        df_contracts['first_5'] = (
                    pd.to_datetime(df_contracts['fixture_date']).dt.year <= 2004)
    df_rig_unique = df_contracts.drop_duplicates(['rig_name', 'first_5'])
    df_rig_unique['n_periods'] = df_rig_unique.groupby(['rig_name'])['first_5'].transform('count')
    df_rig_unique = df_rig_unique[df_rig_unique['n_periods'] == 2]['rig_name'].unique()
    df_contracts = df_contracts[df_contracts['rig_name'].isin(df_rig_unique)]

    #%% GET DEMEANED DAYRATES -----------------------------------------------------------
    df_contracts['month_fixture'] = pd.to_datetime(
        pd.to_datetime(df_contracts['fixture_date']).dt.strftime('%Y-%m')
    )

    df_contracts['day_rate_mean'] = (
        df_contracts
        .groupby('month_fixture')['day_rate']
        .transform('mean')
    )
    df_contracts['day_rate_demean'] = (
        df_contracts['day_rate'] - df_contracts['day_rate_mean']
    )

    #%% TEST RANK CHANGE ----------------------------------------------------------------
    df_contracts['dayrate_mean'] = (
        df_contracts
        .groupby(['first_5', 'rig_name'])['day_rate_demean']
        .transform('mean')
    )
    df_contracts['q_lower'] = (
        df_contracts
        .groupby(['first_5'])['day_rate_demean']
        .transform(lambda x: x.quantile(1/3))
    )
    df_contracts['q_higher'] = (
        df_contracts
        .groupby(['first_5'])['day_rate_demean']
        .transform(lambda x: x.quantile(2/3))
    )
    df_contracts['dayrate_mean'] = (
        df_contracts
        .groupby(['first_5', 'rig_name'])['day_rate_demean']
        .transform('mean')
    )
    df_contracts['price_rank'] = 'none'
    df_contracts.loc[
            df_contracts['dayrate_mean'] >= df_contracts['q_higher'],
            'price_rank'] = 'high'

    df_contracts.loc[(
            (df_contracts['dayrate_mean'] < df_contracts['q_higher'])
            & (df_contracts['dayrate_mean'] > df_contracts['q_lower'])
        ), 'price_rank'] = 'mid'

    df_contracts.loc[
            (df_contracts['dayrate_mean'] <= df_contracts['q_lower']),
            'price_rank'] = 'low'

    #%% GET DF INTO NICE FORMAT ---------------------------------------------------------
    df_rig = (
        df_contracts[['rig_name', 'price_rank', 'first_5']]
        .drop_duplicates()
        .set_index(
            ['rig_name', 'first_5']
        )
        .unstack('first_5')
        .dropna()
    )

    #%% GET PROPORTION SAME RANK --------------------------------------------------------
    df_rig['price_same'] = (
        df_rig[('price_rank', False)] == df_rig[('price_rank', True)]
    )
    output = df_rig['price_same'].mean()

    #%% WRITE OUTPUT --------------------------------------------------------------------
    with open(overleaf_path + f'parameters/parameter_price_rank_{k}.tex', 'w') as fout:
        fout.write(f"{round(output * 100, 1)}\%")
